package org.example;

import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaEmbeddingClient;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;

import java.util.List;
import java.util.Scanner;

public class TestDemo {

    public static void main(String[] args) {
        // 创建Ollama API客户端
        OllamaApi ollamaApi = new OllamaApi();

        // 创建Ollama Embedding客户端，并指定模型
        OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi)
                .withDefaultOptions(OllamaOptions.create().withModel("qwen2.5:7b"));

        // 创建向量存储，用于保存文本嵌入
        VectorStore vectorStore = new SimpleVectorStore(embeddingClient);

        // 添加一些文本到向量存储中
        List<Document> documents = List.of(
                new Document("白日依山尽,黄河入海流。欲穷千里目,更上一层楼。"),
                new Document("青山依旧在,几度夕阳红。白发渔樵江渚上,惯看秋月春风。"),
                new Document("hello。"),
                new Document("你好，世界！")
                // 添加更多文档...
        );
        vectorStore.add(documents);

        // 使用Scanner接收用户输入
        Scanner scanner = new Scanner(System.in);
        while (true){
            System.out.print("请输入关键词进行搜索: ");
            String keyword = scanner.nextLine();

            // 执行相似度搜索
            List<Document> searchResults = vectorStore.similaritySearch(keyword);

            // 输出搜索结果
            System.out.println("查询结果:");
            for (Document doc : searchResults) {
                System.out.println(doc.getContent());
            }
            // 关闭Scanner
            //scanner.close();
        }

    }
}
